#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
节假日及周末公务招待费用分析脚本

此脚本用于分析华北油田在节假日和周末期间发生的公务招待费用情况。
主要功能：
1. 筛选华北油田单位的数据
2. 筛选业务类型为招待的记录
3. 判断开票日期和实际发生日期是否为节假日或周末（非调休工作日）
4. 输出符合条件的数据进行分析
"""

import pandas as pd
from pathlib import Path
from datetime import datetime
import calendar

# 配置常量
DATA_DIR = Path("data")
OUTPUT_DIR = DATA_DIR / "alldata"  # 输出目录
EXPENSE_FILE = DATA_DIR / "费用报销票据信息.xlsx"
HOLIDAY_FILE = DATA_DIR / "节假日信息.xlsx"
OUTPUT_FILE = OUTPUT_DIR / "节假日、周末发生招待分析结果.xlsx"

def is_weekend(date):
    """
    判断日期是否为周末
    
    Args:
        date: datetime对象
        
    Returns:
        bool: 是否为周末
    """
    return date.weekday() >= 5

def check_date_status(date, holidays, workdays):
    """
    检查日期是否符合条件（在节假日列表中或是周末但不是调休工作日）
    
    Args:
        date: datetime对象
        holidays: 节假日集合
        workdays: 调休工作日集合
        
    Returns:
        bool: 是否符合条件
    """
    if pd.isna(date):
        return False
    if date in holidays:
        return True
    if is_weekend(date) and date not in workdays:
        return True
    return False

def load_and_process_data():
    """
    加载并处理数据
    
    Returns:
        tuple: (费用数据DataFrame, 节假日数据DataFrame)
    """
    # 确保输出目录存在
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    # 读取费用数据
    expense_data = pd.read_excel(EXPENSE_FILE)
    # 读取节假日信息
    holiday_data = pd.read_excel(HOLIDAY_FILE)
    
    # 转换日期格式
    expense_data['开票日期'] = pd.to_datetime(expense_data['开票日期']).dt.normalize()  # 只保留年月日
    expense_data['实际发生日期'] = pd.to_datetime(expense_data['实际发生日期']).dt.normalize()  # 只保留年月日
    holiday_data['日期'] = pd.to_datetime(holiday_data['日期']).dt.normalize()  # 只保留年月日
    
    return expense_data, holiday_data

def filter_data(expense_data, holiday_data):
    """
    根据条件筛选数据
    
    Args:
        expense_data: 费用数据DataFrame
        holiday_data: 节假日数据DataFrame
        
    Returns:
        DataFrame: 筛选后的数据
    """
    # 筛选华北油田的数据
    huabei_mask = expense_data['fmis_单位名称'].str.contains('华北油田', na=False)
    
    # 筛选招待相关的数据
    entertainment_mask = expense_data['业务类型名称'].str.contains('招待', na=False)
    
    # 初步筛选
    filtered_data = expense_data[huabei_mask & entertainment_mask].copy()
    
    # 创建节假日集合和调休工作日集合
    holidays = set(holiday_data[holiday_data['类型'] == '节假日']['日期'])
    workdays = set(holiday_data[holiday_data['类型'] == '调休工作日']['日期'])
    
    # 判断开票日期和实际发生日期是否为节假日或周末
    filtered_data['开票日期是否发生招待'] = filtered_data['开票日期'].apply(
        lambda x: '是' if check_date_status(x, holidays, workdays) else '否'
    )
    filtered_data['实际发生日期是否发生招待'] = filtered_data['实际发生日期'].apply(
        lambda x: '是' if check_date_status(x, holidays, workdays) else '否'
    )
    
    # 筛选至少有一个日期为节假日或周末的记录
    result_data = filtered_data[
        (filtered_data['开票日期是否发生招待'] == '是') |
        (filtered_data['实际发生日期是否发生招待'] == '是')
    ].copy()
    
    # 添加日期类型说明
    def get_date_type(date, holidays, workdays):
        if pd.isna(date):
            return '未知'
        if date in holidays:
            return '节假日'
        elif is_weekend(date) and date not in workdays:
            return '周末'
        else:
            return '工作日'
    
    result_data['开票日期类型'] = result_data['开票日期'].apply(
        lambda x: get_date_type(x, holidays, workdays)
    )
    result_data['实际发生日期类型'] = result_data['实际发生日期'].apply(
        lambda x: get_date_type(x, holidays, workdays)
    )
    
    # 格式化日期为年月日格式
    result_data['开票日期'] = result_data['开票日期'].dt.strftime('%Y-%m-%d')
    result_data['实际发生日期'] = result_data['实际发生日期'].dt.strftime('%Y-%m-%d')
    
    return result_data

def main(once = False):
    """主程序入口"""
    try:
        print("开始加载数据...")
        expense_data, holiday_data = load_and_process_data()
        
        print("正在筛选数据...")
        result_data = filter_data(expense_data, holiday_data)
        
        if result_data.empty:
            print("没有找到符合条件的数据")
            return False

        if once:
            return result_data

        # 保存结果
        result_data.to_excel(OUTPUT_FILE, index=False)
        print(f"\n分析完成，结果已保存至: {OUTPUT_FILE}")
        print(f"筛选出的记录数: {len(result_data)}")
        
        # 统计分析
        print("\n数据统计:")
        print(f"仅开票日期为节假日或周末的记录数: {len(result_data[result_data['实际发生日期是否发生招待'] == '否'])}")
        print(f"仅实际发生日期为节假日或周末的记录数: {len(result_data[result_data['开票日期是否发生招待'] == '否'])}")
        print(f"两个日期都为节假日或周末的记录数: {len(result_data[(result_data['开票日期是否发生招待'] == '是') & (result_data['实际发生日期是否发生招待'] == '是')])}")
        
        return True
        
    except Exception as e:
        print(f"处理过程中出现错误: {str(e)}")
        raise

if __name__ == "__main__":
    main() 